#ifndef RKNN_MODEL_RUNNER_H
#define RKNN_MODEL_RUNNER_H

#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include "im2d.h"
#include "RgaUtils.h"
#include "drm_alloc.h"
#include "rknn_api.h"
#include "rk_mpi_mmz.h"

#define ALIGN 8

// 为保持C风格，我们使用结构体定义而非类成员
typedef struct {
    int drm_buffer_fd;
    int drm_buffer_handle;
    size_t actual_size;
    uint8_t *drm_buf;
    uint64_t phy_addr;  
} drm_object_t;

// RKNN模型运行器类
class RknnModelRunner {
public:
    // 构造函数 - 加载模型和分配持久性资源
    RknnModelRunner(const char* model_path);
    
    // 析构函数 - 释放所有资源
    ~RknnModelRunner();
    
    // 从图像文件进行单次推理
    // int inference_from_file(const char* image_path, int loop_count);
    
    // 从摄像头进行单次推理
    int inference_from_camera(int camera_index, int loop_count);
    
    // 从RGBA图像数据进行推理
    int inference_from_image_data(uint8_t* img_data, int img_width, int img_height, int loop_count);
    
    // 从物理连续内存执行RGA缩放
    int do_rga_resize(uint8_t* img_data, int img_width, int img_height);
    
    // 运行推理
    int run_inference(int loop_count);
    
    // 判断模型是否已成功加载
    bool is_model_loaded() const { return ctx != 0; }

private:
    // 禁止拷贝构造和赋值操作
    RknnModelRunner(const RknnModelRunner&);
    RknnModelRunner& operator=(const RknnModelRunner&);

    // 初始化模型和相关资源
    int init_model(const char* model_path);
    
    // 分配输入输出内存
    int allocate_io_memory();
    
    // 释放所有资源
    void free_all_resources();
    
    // 释放临时资源（每次推理后）
    // void free_temp_resources();

    // camera info
    #define CAMERA_WIDTH 640
    #define CAMERA_HEIGHT 480
    #define CAMERA_CHANNELS 3
    #define CAMERA_BUF_SIZE (CAMERA_WIDTH * CAMERA_HEIGHT * CAMERA_CHANNELS)

    // 持久性资源 - 在对象生命周期内保持
    rknn_context ctx;                    // 模型上下文-
    rknn_input_output_num io_num;        // 模型的输入输出信息-
    rknn_tensor_attr* input_attrs;       // 输入张量属性-
    rknn_tensor_attr* output_attrs;      // 输出张量属性-
    int model_width;                     // 模型输入宽度-
    int model_height;                    // 模型输入高度-
    int model_channel;                   // 模型输入通道数-
    
    // 模型输出内存相关
    rknn_tensor_mem** output_mems;          // 输入的RKNN内存-
    
    // 模型输入（rga输出）内存相关
    rga_buffer_t dst;                    // RGA的输出-
    rknn_tensor_mem** input_mems;        // 输入的RKNN内存-

    // rga输入内存相关
    // 预分配的图像缓冲区
    uint8_t img_buf[CAMERA_BUF_SIZE];
    rga_buffer_t src;
    drm_object_t src_drm;                // 源DRM对象-
    rga_buffer_handle_t src_handle;      // 源RGA句柄-

    // 初始化标志
    bool is_initialized;                 // 是否成功初始化-
};

#endif // RKNN_MODEL_RUNNER_H
